#pragma once

#include <bitset>
#include <string>
#include <vector>
#include "../LinkList/LinkListP.h"


struct HuffmanNode
{
	std::string data;
	uint8_t weight;
	std::string code;

	HuffmanNode* parent;
	HuffmanNode* left;
	HuffmanNode* right;

	HuffmanNode(const std::string& inData, uint8_t inWeight):data(inData), weight(inWeight), code(std::string{}), parent(nullptr), left(nullptr), right(nullptr) {}
	HuffmanNode(const std::string& inData):data(inData), weight(1), code(std::string{}), parent(nullptr), left(nullptr), right(nullptr){}

	bool operator==(const HuffmanNode& rItem) const
	{
		return data == rItem.data;
	}

	bool operator<=(const HuffmanNode& rItem) const
	{
		return weight <= rItem.weight;
	}
	
	friend std::ostream& operator<<(std::ostream& os, const HuffmanNode& item)
	{
		os << item.data << ", " << static_cast<unsigned int>(item.weight) << ", " << item.code;

		return os;
	}
};

class HuffmanTree
{
public:
	HuffmanTree(const char* sData);

	// 打印编码表
	void ShowHuffmanCodeTable();
	
	// 可视化树
	void ShowTreeGraph() const;

	// 获取树高
	int Height() const;

	std::string GetCode() const {return DataCode;}
	std::string DecodeData();
private:
	// 构建树
	void BuildTree();

	// 递归计算节点高度
	int GetHeightRecursively(const HuffmanNode* childRoot) const;

	// 填充二叉树可视化缓冲区
	void FillDrawBuffer(HuffmanNode* childRoot, std::vector<std::vector<std::string>>& buffer,
						int row, int col, int hier) const;

	// 对符号进行编码
	void EncodeSymbol();
	// 递归叶子节点
	void RecursiveLeaf(HuffmanNode* childRoot, std::string& codeContainer);
	// 对数据进行编码
	void EncodeData(const char* inData);

	// 递归解码
	void RecursiveDecode(const char*& codeIter, std::string& outSymbol);
private:
	// 根
	HuffmanNode* root;
	// 哈夫曼编码表
	LinkList<HuffmanNode*> HuffmanCodeTable;
	// 哈夫曼编码
	std::string DataCode;
};

HuffmanTree::HuffmanTree(const char* sData)
{
	// 解析文本符号和符号频率
	const char* iter = sData;
	while(*iter != '\0')
	{
		std::string currentChar{*iter};
		HuffmanNode comparedNode{currentChar};
		// 从表中搜索节点
		LinkList<HuffmanNode*>::Node* findedNode = HuffmanCodeTable.Find(&comparedNode, [](const HuffmanNode* a, const HuffmanNode* b)
		{
			if(a->data == b->data)
				return true;

			return false;
		});
		
		if(findedNode)
		{
			// 如果在表中找到了该节点，说明之前已经出现过，所以增加权重
			
			findedNode->data->weight += 1;
		}
		else
		{
			// 如果表中没有找到当前节点，说明这是一个新节点，需要添加到表中
			
			HuffmanNode* huffNode = new HuffmanNode{currentChar};
			HuffmanCodeTable.Insert(huffNode);
		}
		
		++iter;
	}
	// 空表不做处理
	if(HuffmanCodeTable.Size() == 0)
		return;
	
	// 对链表递增排序
	HuffmanCodeTable.Order([](const HuffmanNode* a, const HuffmanNode* b)
		{
			if(a->weight <= b->weight)
				return true;

			return false;
		});
	// 根据已解析的符号链表构建哈夫曼树
	BuildTree();

	// 将哈夫曼编码存入表中
	EncodeSymbol();

	// 对数据进行编码
	EncodeData(sData);
}

void HuffmanTree::ShowHuffmanCodeTable()
{
	LinkList<HuffmanNode*>::Node* tableIter = HuffmanCodeTable.GetHead();
	while(tableIter)
	{
		std::cout << tableIter->data->data << ", " << static_cast<unsigned int>(tableIter->data->weight) << ", " << tableIter->data->code << " -> ";
		tableIter = tableIter->next;
	}
	std::cout << "\n";
}

void HuffmanTree::BuildTree()
{
	int tableSize = HuffmanCodeTable.Size();
	if(tableSize < 2)
	{
		// 说明此时表中只有一个节点，直接设置为根
		root = HuffmanCodeTable.Release(0);
	
		return;
	}

	// 如果表中元素大于等于2，每次取前两个元素
	while(tableSize >= 2)
	{
		HuffmanNode* item0 = HuffmanCodeTable.Release(0);
		HuffmanNode* item1 = HuffmanCodeTable.Release(0);
		
		uint8_t parentWeight = item0->weight + item1->weight;
		HuffmanNode* parent = new HuffmanNode(std::to_string(parentWeight), parentWeight);
		parent->left = item0;
		item0->parent = parent;
		parent->right = item1;
		item1->parent = parent;

		// 将新的节点放回表中
		HuffmanCodeTable.Insert(parent);
		// 重新递增排序
		HuffmanCodeTable.Order([](const HuffmanNode* a, const HuffmanNode* b)
		{
			if(a->weight <= b->weight)
				return true;

			return false;
		});
		tableSize = HuffmanCodeTable.Size();
	}

	root = HuffmanCodeTable.Release(0);
	// ----- 至此完成哈夫曼树的构建 -----
}

#pragma region [print] -------------------------------------------------------------------------------------------------
void HuffmanTree::ShowTreeGraph() const
{
	if(!root)
		return;

	// 矩阵行数 == 树层级数
	int matrixHeight = Height() + 1;
	// 矩阵列数 == (2 * 树最深层的最大节点数) - 1
	int matrixWidth = std::pow(2, matrixHeight) - 1;
	std::vector<std::vector<std::string>> treeMatrix(matrixHeight, std::vector<std::string>(matrixWidth, " "));

	FillDrawBuffer(root, treeMatrix, 0, (matrixWidth - 1)/2, matrixHeight);

	for(const auto& row : treeMatrix)
	{
		for(const auto& item : row)
		{
			std::cout << item ;
		}
		std::cout << std::endl;
	}
	std::cout << "\n" << std::endl;
}

int HuffmanTree::Height() const
{
	if(root)
	{
		// 如果 root 不为空
		int lHeight = GetHeightRecursively(root->left);
		int rHeight = GetHeightRecursively(root->right);

		// 返回两者中大的那个
		// + 根节点高度
		return (lHeight < rHeight ? rHeight : lHeight) + 1;
	}
	else
	{
		// 如果 root 为空，说明树为空，高度不存在
		return -1;
	}
}

int HuffmanTree::GetHeightRecursively(const HuffmanNode* childRoot) const
{
	if(childRoot)
	{
		// 如果 childRoot 不为空
		int lHeight = GetHeightRecursively(childRoot->left);
		int rHeight = GetHeightRecursively(childRoot->right);

		// 返回两者中大的那个
		// + 根节点高度
		return (lHeight < rHeight ? rHeight : lHeight) + 1;
	}
	else
	{
		// 如果 childRoot 为空，说明树为空，高度不存在
		return -1;
	}
}

void HuffmanTree::FillDrawBuffer(HuffmanNode* childRoot, std::vector<std::vector<std::string>>& buffer, int row,
	int col, int hier) const
{
	if(!childRoot)
		return;

	// 填充元素
	buffer[row][col] = childRoot->data;

	// 计算下一层元素列位置偏移
	// hier - row == 当前层数
	// 当前层数 - 1 == 下一层数
	// 层数 - 1 = 当前层的实际树高
	// hier - row - 2 == 下一层的实际树高
	// gap == 下一层子树的最深层的最大节点数
	int gap = std::pow(2, hier - row - 2);

	if(childRoot->left)
		FillDrawBuffer(childRoot->left, buffer, row + 1, col - gap, hier);
	else if( row + 1 < hier)
		buffer[row+1][col - gap] = ".";
	
	if(childRoot->right)
		FillDrawBuffer(childRoot->right, buffer, row + 1, col + gap, hier);
	else if( row + 1 < hier)
		buffer[row+1][col + gap] = ".";
}
#pragma endregion

void HuffmanTree::EncodeSymbol()
{
	std::string codes;
	RecursiveLeaf(root, codes);
	HuffmanCodeTable.Order([](const HuffmanNode* a, const HuffmanNode* b)
		{
			if(a->weight <= b->weight)
				return true;

			return false;
		});
}

void HuffmanTree::RecursiveLeaf(HuffmanNode* childRoot, std::string& codeContainer)
{
	if(!childRoot)
		return;

	if(childRoot->parent)
	{
		// 如果节点是左子节点，编码为0，否则为1
		if(childRoot == childRoot->parent->left)
			codeContainer += "0";
		else
			codeContainer += "1";
	}
	
	// 检测当前节点是否是叶子节点
	if(!(childRoot->left || childRoot->right))
	{
		childRoot->code = codeContainer;
		HuffmanCodeTable.Insert(childRoot);
	}

	RecursiveLeaf(childRoot->left, codeContainer);
	RecursiveLeaf(childRoot->right, codeContainer);
	if(!codeContainer.empty())
		codeContainer.pop_back();
}

void HuffmanTree::EncodeData(const char* inData)
{
	const char* iter = inData;
	// 逐字符编码
	while(*iter != '\0')
	{
		std::string currentChar{*iter};
		HuffmanNode comparedNode{currentChar};
		LinkList<HuffmanNode*>::Node* findedNode = HuffmanCodeTable.Find(&comparedNode, [](const HuffmanNode* a, const HuffmanNode* b)
		{
			if(a->data == b->data)
				return true;

			return false;
		});

		DataCode += findedNode->data->code;
		
		++iter;
	}
}

std::string HuffmanTree::DecodeData()
{
	if(DataCode.empty())
		return "";

	std::string data;
	// 获取编码
	const char* codeIter = DataCode.c_str();
	// 逐字符解码
	while(*codeIter != '\0')
	{
		std::string symbol;
		RecursiveDecode(codeIter, symbol);

		data += symbol;
	}
	
	return data;
}

void HuffmanTree::RecursiveDecode(const char*& codeIter, std::string& outSymbol)
{
	HuffmanNode*  iterNode = root;
	while(*codeIter != '\0')
	{
		// 根据编码符号是 0 还是 1 来决定路径走左边还是右边
		char current = *codeIter;
		if(current == '0')
			iterNode = iterNode->left;
		else
			iterNode = iterNode->right;

		++codeIter;
		
		// 如果当前节点是叶子节点，说明完成一次解码
		if(!(iterNode->left || iterNode->right))
		{
			outSymbol = iterNode->data;
			break;
		}
	}
}